from tqdm import tqdm
import math
import numpy as np
import ROOT as rt

rt.gROOT.ProcessLine("struct TreeVar_t {\
    Float_t         Energy;\
    Float_t         Weight;\
};" )

nMCBins = int(40)
nDalitzBins = int(200)
nDalitzMeerkatBins = int(2000)
nDivBins = int(20)

kernelWidth = 0.005
sampleBins = 1500

energyRange = [0., 10.2]

def GenRoot():
    TreeVar = rt.TreeVar_t()
    InTree = rt.TChain("MCData")
    InTree.Add("GammaInDME.root")
    InTree.Add("GammaInDMEEnlarge.root")
    NumOfEntries = InTree.GetEntries()

    OutFile = rt.TFile("GammaEnergyFloat.root", "recreate")
    OutTree = rt.TTree("MCData", "")
    OutTree.Branch("Energy", rt.addressof(TreeVar, "Energy"), "Energy/F")
    OutTree.Branch("Weight", rt.addressof(TreeVar, "Weight"), "Weight/F")

    for entry in tqdm(InTree, total=NumOfEntries):
        TreeVar.Energy = entry.PassEng
        TreeVar.Weight = 1.
        OutTree.Fill()

    OutFile.cd()
    OutTree.AutoSave()
    OutFile.Close()

def GenMap():

    import os, sys
    os.environ["ROOT_INCLUDE_PATH"] = "$MEERKATROOT/inc/"
    rt.gSystem.Load("$MEERKATROOT/lib/libMeerkat.so")

    InFile = rt.TFile(f"GammaEnergyFloat.root")
    InTree = InFile.Get("MCData")

    energyPHSP  = rt.OneDimPhaseSpace("energyPHSP" , energyRange[0], energyRange[1])

    GammaEnergy = rt.TH1F("GammaEnergy", "Dalitz plot", nDalitzMeerkatBins, energyRange[0], energyRange[1])
    GammaEnergy.GetXaxis().SetTitle("Theta")
    GammaEnergy.GetYaxis().SetTitle("Phi")

    KDE_gen = rt.BinnedKernelDensity(
            "KernelPDF",
            energyPHSP,
            InTree,
            "Energy",
            "Weight",
            sampleBins,
            kernelWidth,
            0,
            20000000
        )

    OutFile = rt.TFile("GammaEnergy.root", "recreate")
    OutFile.cd()
    KDE_gen.project(GammaEnergy)

    GammaEnergy.Write("GammaEnergy")

    GammaEnergyTruth = rt.TH1F("GammaEnergyTruth", "Dalitz plot", nDalitzMeerkatBins, energyRange[0], energyRange[1])
    InTree.Project("GammaEnergyTruth", "Energy", '1')
    GammaEnergyTruth.Write()

    #InFile.Close()
    #OutFile.Close()

if __name__ == "__main__":
        #GenRoot()
        GenMap()
